"""
实体记忆管理器

功能：
- 记录债券、公司等实体的关键信息
- 支持实体信息的增量更新
- 提供实体搜索功能
- 支持数据持久化

作者：System
创建时间：2025-01-13
"""

from typing import List, Dict, Optional
from datetime import datetime
import json
import os
from pathlib import Path


class EntityMemory:
    """实体记忆管理器"""
    
    def __init__(self, storage_path: str = "cache/entity_memory.json"):
        """
        初始化实体记忆管理器
        
        Args:
            storage_path: 实体信息存储路径
        """
        self.storage_path = storage_path
        self.entities = {}  # {entity_id: entity_info}
        self.load_entities()
    
    def add_or_update_bond(self, bond_code: str, info: dict) -> None:
        """
        添加或更新债券信息
        
        Args:
            bond_code: 债券代码
            info: 债券信息字典
        """
        entity_id = f"bond_{bond_code}"
        
        if entity_id not in self.entities:
            self.entities[entity_id] = {
                "type": "bond",
                "code": bond_code,
                "info": {},
                "last_updated": None,
                "created_at": datetime.now().isoformat()
            }
        
        # 增量更新信息
        self.entities[entity_id]["info"].update(info)
        self.entities[entity_id]["last_updated"] = datetime.now().isoformat()
        
        # 自动保存
        self.save_entities()
    
    def add_or_update_company(self, company_name: str, info: dict) -> None:
        """
        添加或更新公司信息
        
        Args:
            company_name: 公司名称
            info: 公司信息字典
        """
        entity_id = f"company_{company_name}"
        
        if entity_id not in self.entities:
            self.entities[entity_id] = {
                "type": "company",
                "name": company_name,
                "info": {},
                "last_updated": None,
                "created_at": datetime.now().isoformat()
            }
        
        # 增量更新信息
        self.entities[entity_id]["info"].update(info)
        self.entities[entity_id]["last_updated"] = datetime.now().isoformat()
        
        # 自动保存
        self.save_entities()
    
    def add_or_update_entity(self, entity_type: str, entity_key: str, info: dict) -> None:
        """
        通用的添加或更新实体方法
        
        Args:
            entity_type: 实体类型（如 "bond", "company", "industry"）
            entity_key: 实体唯一标识（如债券代码、公司名称）
            info: 实体信息字典
        """
        entity_id = f"{entity_type}_{entity_key}"
        
        if entity_id not in self.entities:
            self.entities[entity_id] = {
                "type": entity_type,
                "key": entity_key,
                "info": {},
                "last_updated": None,
                "created_at": datetime.now().isoformat()
            }
        
        # 增量更新信息
        self.entities[entity_id]["info"].update(info)
        self.entities[entity_id]["last_updated"] = datetime.now().isoformat()
        
        # 自动保存
        self.save_entities()
    
    def get_bond_info(self, bond_code: str) -> Optional[dict]:
        """
        获取债券信息
        
        Args:
            bond_code: 债券代码
            
        Returns:
            债券信息字典，如果不存在则返回None
        """
        entity_id = f"bond_{bond_code}"
        return self.entities.get(entity_id)
    
    def get_company_info(self, company_name: str) -> Optional[dict]:
        """
        获取公司信息
        
        Args:
            company_name: 公司名称
            
        Returns:
            公司信息字典，如果不存在则返回None
        """
        entity_id = f"company_{company_name}"
        return self.entities.get(entity_id)
    
    def get_entity(self, entity_type: str, entity_key: str) -> Optional[dict]:
        """
        通用的获取实体方法
        
        Args:
            entity_type: 实体类型
            entity_key: 实体唯一标识
            
        Returns:
            实体信息字典，如果不存在则返回None
        """
        entity_id = f"{entity_type}_{entity_key}"
        return self.entities.get(entity_id)
    
    def search_entities(self, keyword: str, entity_type: Optional[str] = None, limit: int = 20) -> List[dict]:
        """
        搜索实体
        
        Args:
            keyword: 搜索关键词
            entity_type: 实体类型（可选，不指定则搜索所有类型）
            limit: 最多返回结果数（默认20）
            
        Returns:
            匹配的实体列表
        """
        results = []
        keyword_lower = keyword.lower()
        
        for entity_id, entity in self.entities.items():
            # 如果指定了实体类型，过滤掉不匹配的
            if entity_type and entity["type"] != entity_type:
                continue
            
            # 简单的关键词匹配（在实体ID、类型、信息中搜索）
            entity_str = json.dumps(entity, ensure_ascii=False).lower()
            if keyword_lower in entity_str or keyword_lower in entity_id.lower():
                results.append(entity)
            
            # 限制返回数量
            if len(results) >= limit:
                break
        
        return results
    
    def get_all_bonds(self) -> List[dict]:
        """
        获取所有债券实体
        
        Returns:
            所有债券实体列表
        """
        return [
            entity for entity_id, entity in self.entities.items()
            if entity["type"] == "bond"
        ]
    
    def get_all_companies(self) -> List[dict]:
        """
        获取所有公司实体
        
        Returns:
            所有公司实体列表
        """
        return [
            entity for entity_id, entity in self.entities.items()
            if entity["type"] == "company"
        ]
    
    def get_entities_by_type(self, entity_type: str) -> List[dict]:
        """
        按类型获取实体
        
        Args:
            entity_type: 实体类型
            
        Returns:
            指定类型的所有实体列表
        """
        return [
            entity for entity in self.entities.values()
            if entity["type"] == entity_type
        ]
    
    def delete_entity(self, entity_type: str, entity_key: str) -> bool:
        """
        删除实体
        
        Args:
            entity_type: 实体类型
            entity_key: 实体唯一标识
            
        Returns:
            是否删除成功
        """
        entity_id = f"{entity_type}_{entity_key}"
        if entity_id in self.entities:
            del self.entities[entity_id]
            self.save_entities()
            return True
        return False
    
    def delete_bond(self, bond_code: str) -> bool:
        """
        删除债券实体
        
        Args:
            bond_code: 债券代码
            
        Returns:
            是否删除成功
        """
        return self.delete_entity("bond", bond_code)
    
    def delete_company(self, company_name: str) -> bool:
        """
        删除公司实体
        
        Args:
            company_name: 公司名称
            
        Returns:
            是否删除成功
        """
        return self.delete_entity("company", company_name)
    
    def clear_old_entities(self, days: int = 90) -> int:
        """
        清理长期未更新的实体
        
        Args:
            days: 保留最近N天更新的实体（默认90天）
            
        Returns:
            清理的实体数
        """
        from datetime import timedelta
        
        cutoff_time = datetime.now() - timedelta(days=days)
        cleared_count = 0
        
        for entity_id in list(self.entities.keys()):
            entity = self.entities[entity_id]
            last_updated = entity.get("last_updated")
            
            # 如果从未更新过，检查创建时间
            if not last_updated:
                last_updated = entity.get("created_at")
            
            if last_updated:
                update_time = datetime.fromisoformat(last_updated)
                if update_time < cutoff_time:
                    del self.entities[entity_id]
                    cleared_count += 1
        
        if cleared_count > 0:
            self.save_entities()
        
        return cleared_count
    
    def save_entities(self) -> None:
        """保存实体到文件"""
        try:
            # 确保目录存在
            Path(self.storage_path).parent.mkdir(parents=True, exist_ok=True)
            
            with open(self.storage_path, 'w', encoding='utf-8') as f:
                json.dump(self.entities, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"保存实体信息失败: {e}")
    
    def load_entities(self) -> None:
        """从文件加载实体"""
        try:
            if os.path.exists(self.storage_path):
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    self.entities = json.load(f)
        except FileNotFoundError:
            pass
        except Exception as e:
            print(f"加载实体信息失败: {e}")
    
    def get_statistics(self) -> Dict:
        """
        获取统计信息
        
        Returns:
            包含统计信息的字典
        """
        type_counts = {}
        for entity in self.entities.values():
            entity_type = entity["type"]
            type_counts[entity_type] = type_counts.get(entity_type, 0) + 1
        
        return {
            "total_entities": len(self.entities),
            "entity_types": type_counts,
            "bonds_count": type_counts.get("bond", 0),
            "companies_count": type_counts.get("company", 0)
        }


# 创建全局单例（可选）
_global_entity_memory = None

def get_entity_memory() -> EntityMemory:
    """
    获取全局实体记忆管理器单例
    
    Returns:
        EntityMemory实例
    """
    global _global_entity_memory
    if _global_entity_memory is None:
        _global_entity_memory = EntityMemory()
    return _global_entity_memory